Skip to content

Cosmos3 super reasoner - #5074

Merged
copybara-service[bot] merged 1 commit into
mainfrom
cosmos3-super-reasoner
Sep 4, 2026
Merged

Cosmos3 super reasoner#5074
copybara-service[bot] merged 1 commit into
mainfrom
cosmos3-super-reasoner

Conversation

@lydhr

@lydhr lydhr commented Aug 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

This PR adds support for onboarding and running inference on the NVIDIA cosmos3-super, reasoner tower.

Key updates

  • Reuse: It leverages the existing architectural blocks derived from Qwen3/Qwen3-VL, sharing preprocessing and attention mechanisms with the previously implemented Cosmos3-Nano variants while scaling up to a larger qwen3-vl backbone.

  • New components:

    • Dimensions: $d_{model}=5120$, $L_{heads}=64$, $L_{kv_heads}=8$.
    • Capacity: 64 decoder layers.

  • Fix the processor for cosmos models: the identifier should be the vision_block name rather than model name; thus the redundant explicit listing of cosmos3-xxx are removed.

Tests

Checkpoint conversion

MODEL_NAME=cosmos3-super-reasoner

PYTHONPATH=src/ python3 -m maxtext.checkpoint_conversion.to_maxtext \
    src/maxtext/configs/base.yml \
    model_name=${MODEL_NAME?} \
    base_output_directory=<Your GS bucket or local>\
    scan_layers=false \
    use_multimodal=true \
    hardware=cpu \
    skip_jax_distributed_system=True \
    checkpoint_storage_use_ocdbt=False \
    checkpoint_storage_use_zarr3=False \
    --eager_load_method=safetensors \
    --lazy_load_tensors=False

Decoding (Robotic Arm Planning) (v6 lite)

MODEL_NAME=cosmos3-super-reasoner
export checkpoint_path=gs://ly-maxtext/output/cosmos3-super-reasoner/09031507/0/items

PYTHONPATH=src/ python3 -m maxtext.inference.decode \
    src/maxtext/configs/base.yml \
    model_name=${MODEL_NAME?} \
    tokenizer_path=nvidia/Cosmos3-Super \
    tokenizer_type=huggingface  \
    load_parameters_path=${checkpoint_path} \
    per_device_batch_size=1 \
    run_name=${MODEL_NAME?}_decode \
    scan_layers=false \
    use_multimodal=true \
    prompt='You are a robotic arm planner. The task is to put the flower into the red bottle. Generate a detailed plan consisting of sequential subtasks to accomplish the task.' \
    image_path='tests/assets/test_image_reasoning.jpg' \
    max_prefill_predict_length=1024 \
    max_target_length=2248 \. # increased the length to see the full reasoning result.
    ici_tensor_parallelism=4 \
    override_model_config=true \
    attention='dot_product'
image
Decoding output and analysis: Input `<|im_start|>user <|vision_start|><|image_pad|><|vision_end|>You are a robotic arm planner. The task is to put the flower into the red bottle. Generate a detailed plan consisting of sequential subtasks to accomplish the task.<|im_end|> <|im_start|>assistant ` -> `

Task Analysis:

The goal is to place the flower into the red bottle. To achieve this, we need to break down the task into a series of logical steps that the robotic arm can execute sequentially. Below is a detailed plan with explanations for each step.

Step 1: Identify and Locate the Flower

  • Objective: Determine the position of the flower on the table.
  • Explanation: The flower is currently lying flat on the table. The robotic arm must first locate it using its sensors or vision system. This step ensures that the arm knows where to move to pick up the flower.

Step 2: Move the Robotic Arm to the Flower

  • Objective: Position the robotic arm's gripper over the flower.
  • Explanation: Once the flower's location is identified, the robotic arm needs to move its gripper to the exact coordinates of the flower. This involves precise movement along the X, Y, and Z axes to align the gripper directly above the flower.

Step 3: Grasp the Flower

  • Objective: Securely pick up the flower using the gripper.
  • Explanation: The gripper must close around the flower's stem to ensure a firm grip. It is important to avoid damaging the flower during this process. The gripper should be designed to handle delicate objects like flowers.

Step 4: Lift the Flower

  • Objective: Raise the flower off the table.
  • Explanation: After grasping the flower, the robotic arm lifts it vertically to clear the table surface. This step ensures that the flower is fully detached from the table and ready to be moved toward the red bottle.

Step 5: Identify and Locate the Red Bottle

  • Objective: Determine the position of the red bottle on the table.
  • Explanation: The red bottle is one of the three bottles on the table. The robotic arm must use its sensors or vision system to identify which bottle is red and note its coordinates.

Step 6: Move the Robotic Arm to the Red Bottle

  • Objective: Position the robotic arm's gripper over the red bottle.
  • Explanation: With the flower securely held, the robotic arm moves horizontally and vertically to bring the flower above the opening of the red bottle. Precision is crucial here to ensure the flower is aligned correctly for insertion.

Step 7: Insert the Flower into the Red Bottle

  • Objective: Place the flower inside the red bottle.
  • Explanation: The robotic arm carefully lowers the flower into the red bottle. The gripper should release the flower once it is fully inside the bottle. This step requires gentle control to avoid knocking over the bottle or damaging the flower.

Step 8: Release the Flower

  • Objective: Ensure the flower is securely placed inside the red bottle.
  • Explanation: After inserting the flower, the gripper opens to release it. The robotic arm then retracts slightly to confirm that the flower is stable within the bottle.

Step 9: Verify the Task Completion

  • Objective: Confirm that the flower is successfully placed in the red bottle.
  • Explanation: The robotic arm may use its sensors or vision system to verify that the flower is indeed inside the red bottle. If any issues are detected (e.g., the flower is not properly placed), the arm can repeat the necessary steps.

Final Answer:

The detailed plan to put the flower into the red bottle is as follows:

  1. Identify and Locate the Flower
  2. Move the Robotic Arm to the Flower
  3. Grasp the Flower
  4. Lift the Flower
  5. Identify and Locate the Red Bottle
  6. Move the Robotic Arm to the Red Bottle
  7. Insert the Flower into the Red Bottle
  8. Release the Flower
  9. Verify the Task Completion

By following these steps sequentially, the robotic arm will successfully complete the task of placing the flower into the red bottle.

  1. Identify and Locate the Flower
  2. Move the Robotic Arm to the Flower
  3. Grasp the Flower
  4. Lift the Flower
  5. Identify and Locate the Red Bottle
  6. Move the Robotic Arm to the Red Bottle
  7. Insert the Flower into the Red Bottle
  8. Release the Flower
  9. Verify the Task Completion`

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

gemini-code-assist[bot]
gemini-code-assist Bot reviewed Aug 31, 2026
View reviewed changes

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request adds support for NVIDIA Cosmos3-Nano and Cosmos3-Super Reasoner models, including their configurations, parameter mappings, and multimodal processing logic. It also introduces output truncation upon encountering an EOS token during decoding and switches safetensor loading to recursive search. The review feedback highlights a potential bug in handling list-valued eos_token_ids, a configuration mismatch regarding tied word embeddings for the Cosmos3-Nano model, and risks associated with recursively loading safetensors from subdirectories.

Comment thread src/maxtext/inference/decode.py
Comment thread src/maxtext/checkpoint_conversion/utils/utils.py
@lydhr
lydhr force-pushed the cosmos3-super-reasoner branch 2 times, most recently from 16453a7 to 161e435 Compare September 3, 2026 22:53
@lydhr
lydhr force-pushed the cosmos3-super-reasoner branch from 161e435 to 00aec75 Compare September 3, 2026 23:19
@AI-Hypercomputer AI-Hypercomputer deleted a comment from gemini-code-assist Bot Sep 3, 2026
@lydhr
lydhr marked this pull request as ready for review September 3, 2026 23:26
@codecov

codecov Bot commented Sep 3, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@copybara-service
copybara-service Bot merged commit 4cfec5b into main Sep 4, 2026
98 of 99 checks passed
@copybara-service
copybara-service Bot deleted the cosmos3-super-reasoner branch September 4, 2026 21:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@bvandermoon bvandermoon bvandermoon approved these changes
@YixuanWang-99 YixuanWang-99 YixuanWang-99 approved these changes
@parambole parambole Awaiting requested review from parambole parambole is a code owner
@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin is a code owner
@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner
@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner
@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner
@gagika gagika Awaiting requested review from gagika gagika is a code owner
@shralex shralex Awaiting requested review from shralex shralex is a code owner
@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner
@NuojCheng NuojCheng Awaiting requested review from NuojCheng NuojCheng is a code owner
@xibinliu xibinliu Awaiting requested review from xibinliu xibinliu is a code owner
@huytransformer huytransformer Awaiting requested review from huytransformer huytransformer is a code owner
@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner
@igorts-git igorts-git Awaiting requested review from igorts-git igorts-git is a code owner
@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner
@abhinavclemson abhinavclemson Awaiting requested review from abhinavclemson abhinavclemson is a code owner
@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner
@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner
@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner
@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner
@dipannita08 dipannita08 Awaiting requested review from dipannita08 dipannita08 is a code owner
@darisoy darisoy Awaiting requested review from darisoy darisoy is a code owner
+1 more reviewer
@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@lydhr @bvandermoon @YixuanWang-99